from concurrent.futures import process
import json
import os
from multiprocessing import Process, Queue
import re

cpt_dir = os.path.join(
    "/nfs/home/share/jiaxiaoyu/simpoint_checkpoint_archive/"
    "spec17-rv64gcb-O3-20m-gcc12.2.0-mix-with-special_wrf", 
    "checkpoint-0-0-0"
)
cpt_data = os.listdir(cpt_dir)
zstd_pattern = re.compile(r"(.*)_(.*)_(.*)_.zstd")
total_cases = 0

write_buffer = {}

for case in cpt_data:
  case_dir = os.path.join(cpt_dir, case)

  if not os.path.isdir(case_dir):
    continue

  total_cases += 1
  points = os.listdir(case_dir)
  points_buffer = {}
  for point in points:
    point_dir = os.path.join(case_dir, point)
    cpt = os.listdir(point_dir)
    assert(len(cpt) <= 1)
    zstd_match = zstd_pattern.match(str(cpt))
    if zstd_match is not None:
      points_buffer.update({str(zstd_match.group(2)) : str(zstd_match.group(3))})

  # case_buffer = {"points" : points_buffer}
  write_buffer.update({str(case) : points_buffer})

print("total cases: ", total_cases)
with open('summary.json', 'w+') as f:
  json.dump(write_buffer, f, indent=4)



